using System;
using System.Collections.Generic;
using System.Runtime.InteropServices;

namespace NTrace
{
  internal class TraceMessageBuilder : IDisposable
  {
    private static readonly int sizeofEventTraceHeader = Marshal.SizeOf(typeof(EventTraceHeader));
    private static readonly int sizeofMofField = Marshal.SizeOf(typeof(MofField));

    /// <summary>
    /// A pointer to the unmanaged block of memory that contains our trace message
    /// </summary>
    private IntPtr message;

    private IList<IntPtr> stringPointers;

    /// <summary>
    /// The pointer for writing our structures
    /// </summary>
    private IntPtr writePosition;

    /// <summary>
    /// The pointer to the data block for our message
    /// </summary>
    private IntPtr dataPointer;

    private Guid messageGuid;

    private int messageNumber;

    /// <summary>
    /// A list of strings to pass as arguments
    /// </summary>
    private IList<String> arguments;
    
    public TraceMessageBuilder()
    {
      this.arguments = new List<String>();
      this.stringPointers = new List<IntPtr>();
    }

    public Guid MessageGuid
    {
      get
      {
        return this.messageGuid;
      }
      set
      {
        this.messageGuid = value;
      }
    }

    public int MessageNumber
    {
      get
      {
        return this.messageNumber;
      }
      set
      {
        this.messageNumber = value;
      }
    }

    public void AddStringArgument(String argument)
    {
      this.arguments.Add(argument);
    }

    public IntPtr ToPointer()
    {
      // We'll represent all strings as ItemPWString, which is comprised of two MOF_FIELD entries:
      // The first one is a ushort containing the length of the string, the second one is 
      // a pointer to the string
      //
      // We'll pass (n * 2) + 3 MOF_FIELDs. One for the message number, one for the hash code, n * 2
      // for string lengths plus strings, and another for the terminator
      //
      // The buffer is laid out thusly:
      // <start>
      //   [EventTraceHeader]
      //   [MOF_FIELD - GetHashCode]
      //   [MOF_FIELD - String 1 length]
      //   [MOF_FIELD - String 1 pointer]
      //   [MOF_FIELD - String 2 length]
      //   [MOF_FIELD - String 2 pointer]
      //   [additional MOF_FIELDs as needed]
      //   [NULL MOF_FIELD]
      // <data area>
      //   [ushort for message number]
      //   [DWORD for hash code]
      //   [ushort - string length 1]
      //   [ushort - string length 2]
      //   [additional ushorts for string length]
      // <end>

      // Tell the runtime we should be finalized if someone forgets to dispose us
      GC.ReRegisterForFinalize(this);

      int totalMofFields = (this.arguments.Count * 2) + 2;
      int bufferSize = TraceMessageBuilder.sizeofEventTraceHeader + // Our EventTraceHeader
        TraceMessageBuilder.sizeofMofField * totalMofFields +  // The number of MofField structures we'll be passing
        sizeof(ushort) + sizeof(ushort) * this.arguments.Count; // The size of the data area: message Number (ushort) + ushorts for string lengths

      // The header for our event
      EventTraceHeader eventTraceHeader = new EventTraceHeader();
      eventTraceHeader.Size = (ushort)bufferSize;
      eventTraceHeader.Guid = this.messageGuid;
      eventTraceHeader.Version = 255; // This is undocumented and indicates that we are supplying a message number

      //eventTraceHeader.ClassLevel = level;
      eventTraceHeader.Flags = TraceHeaderFlags.UseGuid | TraceHeaderFlags.DataAttached;

      // Allocate memory to hold our message
      this.message = Marshal.AllocHGlobal(bufferSize);

      // Position our write point at the start of the block
      this.writePosition = message;

      // Position the data pointer just past the MOF_FIELD structs (which come after the event header)
      this.dataPointer = NativeMethods.MovePointer(
        this.message, 
        TraceMessageBuilder.sizeofEventTraceHeader + TraceMessageBuilder.sizeofMofField * totalMofFields);

      // Write our event header to the buffer
      this.WriteHeader(eventTraceHeader);

      // Write the message number field
      this.WriteMessageNumber();

      for (int index = 0; index < this.arguments.Count; index++)
      {
        this.WriteString(arguments[index]);
      }

      this.WriteTerminator();

      return this.message;
    }

    private void WriteTerminator()
    {
      MofField currentField = new MofField();
      currentField.Length = 0;
      currentField.DataPtr = 0;
      Marshal.StructureToPtr(currentField, this.writePosition, false);
    }

    private void WriteString(String argument)
    {
      IntPtr stringPointer = Marshal.StringToHGlobalUni(argument);
      // Add it to the list for disposal
      this.stringPointers.Add(stringPointer);

      ushort stringLengthInBytes = 0;
      if (!String.IsNullOrEmpty(argument))
      {
        stringLengthInBytes = (ushort)System.Text.Encoding.Unicode.GetByteCount(argument);
      }

      // First we'll write the length field, then a pointer to the string
      this.WriteUShort(stringLengthInBytes);

      // Now add the pointer to the string
      MofField currentField = new MofField();
      currentField.Length = (uint)stringLengthInBytes;
      currentField.DataPtr = (UInt64)stringPointer;
      Marshal.StructureToPtr(currentField, this.writePosition, false);

      // Move to the next entry
      this.MoveWritePointer(TraceMessageBuilder.sizeofMofField);
    }

    private void WriteMessageNumber()
    {
      this.WriteUShort((ushort)this.messageNumber);
    }

    /*
    We could have used a generic method for the below (and still may eventually)
    but the sizeof() operator does not work with generics. This would mean that
    we have to call Marshal.SizeOf() to dynamically determine the size for every
    call rather than determining it once at compile time. :(
    */

    private void WriteUShort(ushort value)
    {
      MofField currentField = new MofField();
      currentField.Length = sizeof(ushort);
      // Point to the address at which we will write our ushort
      currentField.DataPtr = (UInt64)this.dataPointer;
      Marshal.StructureToPtr(currentField, this.writePosition, false);
      // Move the write position on to the next entry
      this.MoveWritePointer(TraceMessageBuilder.sizeofMofField);

      // Write the value to the data block at the data pointer's current position
      Marshal.StructureToPtr(value, this.dataPointer, false);
      // Move the data pointer past the number
      this.MoveDataPointer(sizeof(ushort));
    }

    private void WriteInt32(Int32 value)
    {
      MofField currentField = new MofField();
      currentField.Length = sizeof(Int32);
      // Point to the address at which we will write our ushort
      currentField.DataPtr = (UInt64)this.dataPointer;
      Marshal.StructureToPtr(currentField, this.writePosition, false);
      // Move the write position on to the next entry
      this.MoveWritePointer(TraceMessageBuilder.sizeofMofField);

      // Write the value to the data block at the data pointer's current position
      Marshal.StructureToPtr(value, this.dataPointer, false);
      // Move the data pointer past the number
      this.MoveDataPointer(sizeof(Int32));
    }

    private void MoveWritePointer(int amount)
    {
      this.writePosition = NativeMethods.MovePointer(this.writePosition, amount);
    }

    private void MoveDataPointer(int amount)
    {
      this.dataPointer = NativeMethods.MovePointer(this.dataPointer, amount);
    }

    private void WriteHeader(EventTraceHeader eventTraceHeader)
    {
      Marshal.StructureToPtr(eventTraceHeader, this.writePosition, false);
      this.MoveWritePointer(TraceMessageBuilder.sizeofEventTraceHeader);
    }

    ~TraceMessageBuilder()
    {
      this.FreeMemory();
    }

    public void Dispose()
    {
      if (this.message != IntPtr.Zero)
      {
        GC.SuppressFinalize(this);
        this.FreeMemory();
      }
    }

    private void FreeMemory()
    {
      for (int index = 0; index < this.stringPointers.Count; index++)
      {
        if (this.stringPointers[index] != IntPtr.Zero)
        {
          Marshal.FreeHGlobal(this.stringPointers[index]);
        }
      }
      this.stringPointers.Clear();
      Marshal.FreeHGlobal(this.message);
      this.message = IntPtr.Zero;
    }
  }
}
